#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2025-07-18


"""
【节点名称】：
    YOLOv11PoseNode_Cuda
【订阅类型】：
    sensor_msgs::CompressedImage / memory_msgs::RawImage （输入图像）
【发布类型】：
    spirecv_msgs::2DTargets （检测结果）
    sensor_msgs::CompressedImage （可视化结果，需借助可视化工具）
    std_msgs::Boolean （如果输入节点是数据集，则检测完成发布该话题让输入节点继续工作）
"""

from ultralytics import YOLO
import threading
from queue import Queue
import numpy as np
import time
import os
from typing import Union
from spirems import Subscriber, Publisher, cvimg2sms, sms2cvimg, def_msg, QoS, BaseNode, get_extra_args
from spirems.mod_helper import download_model
import argparse
import platform


def trans_det_results(det_results, kpt_links):
    sms_results = def_msg('spirecv_msgs::2DTargets')

    sms_results["file_name"] = ""
    sms_results["height"] = int(det_results.orig_shape[0])
    sms_results["width"] = int(det_results.orig_shape[1])
    sms_results["kpt_links"] = kpt_links
    sms_results["targets"] = []
    cls = det_results.boxes.cls.cpu().numpy().astype(np.float64)
    conf = det_results.boxes.conf.cpu().numpy().astype(np.float64)
    xywh = det_results.boxes.xywh.cpu().numpy().astype(np.float64)
    xyn = det_results.keypoints.xyn.cpu().numpy().astype(np.float64)
    for i in range(len(cls)):
        ann = dict()
        ann["category_name"] = det_results.names[int(cls[i])].strip().replace(' ', '_').lower()
        ann["category_id"] = int(cls[i])
        ann["score"] = round(conf[i], 3)
        ann["bbox"] = [round(j, 3) for j in xywh[i].tolist()]
        ann["bbox"][0] = ann["bbox"][0] - ann["bbox"][2] / 2
        ann["bbox"][1] = ann["bbox"][1] - ann["bbox"][3] / 2
        ann["keypoints"] = [[round(j[0] * sms_results["width"], 3), round(j[1] * sms_results["height"], 3)] for j in xyn[i]]
        sms_results["targets"].append(ann)

    return sms_results


class YOLOv11PoseNode_Cuda(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        threading.Thread.__init__(self)
        sms_shutdown = True if sms_shutdown in ['True', 'true', '1', True] else False
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )
        self.launch_next_emit = self.get_param("launch_next_emit", True)
        self.specified_input_topic = self.get_param("specified_input_topic", "")
        self.specified_output_topic = self.get_param("specified_output_topic", "")
        self.realtime_det = self.get_param("realtime_det", False)
        self.remote_ip = self.get_param("remote_ip", "127.0.0.1")
        self.remote_port = self.get_param("remote_port", 9094)
        self.confidence = self.get_param("confidence", 0.001)
        self.nms_thresh = self.get_param("nms_thresh", 0.60)
        self.imgsz = self.get_param("imgsz", 640)
        self.dataset_name = self.get_param("dataset_name", "coco_pose")
        self.kpt_links = self.get_param("kpt_links", [[5, 7], [7, 9], [6, 8], [8, 10], [11, 13], [13, 15], [12, 14], [14, 16], [11, 12], [5, 6]])
        self.pt_model = self.get_param("pt_model", "sms::yolo11m-pose.pt")
        self.use_shm = self.get_param("use_shm", -1)
        self.params_help()

        self.b_use_shm = False
        if self.use_shm == 1 or (self.use_shm == -1 and platform.system() == 'Linux'):
            self.b_use_shm = True

        if self.pt_model.startswith("sms::"):
            self.local_pt_model = download_model(self.__class__.__name__, self.pt_model)
            assert self.local_pt_model is not None
        else:
            self.local_pt_model = self.pt_model

        input_url = '/' + job_name + '/sensor/image_raw'
        if len(self.specified_input_topic) > 0:
            input_url = self.specified_input_topic

        output_url = '/' + job_name + '/detector/results'
        if len(self.specified_output_topic) > 0:
            output_url = self.specified_output_topic

        self.job_queue = Queue()
        self.queue_pool.append(self.job_queue)

        self._image_reader = Subscriber(
            input_url, 'std_msgs::Null', self.image_callback,
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._result_writer = Publisher(
            output_url, 'spirecv_msgs::2DTargets',
            ip=self.remote_ip, port=self.remote_port, qos=QoS.Reliability
        )
        self._show_writer = Publisher(
            '/' + job_name + '/detector/image_results', 'memory_msgs::RawImage' if self.b_use_shm else 'sensor_msgs::CompressedImage',
            ip=ip, port=port
        )
        if self.launch_next_emit:
            self._next_writer = Publisher(
                '/' + job_name + '/launch_next', 'std_msgs::Boolean',
                ip=ip, port=port, qos=QoS.Reliability
            )

        self._detector = YOLO(self.local_pt_model)  # load an official model
        self.infer_n_ims = 0
        self.infer_time3 = 0.0
        self.infer_time2 = 0.0
        self.infer_time1 = 0.0

        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_reader.kill()
        self._result_writer.kill()
        self._show_writer.kill()
        self._next_writer.kill()
        self._image_reader.join()
        self._result_writer.join()
        self._show_writer.join()
        self._next_writer.join()

    def image_callback(self, msg):
        if self.realtime_det:
            if not self.job_queue.empty():
                self.job_queue.queue.clear()
        img = sms2cvimg(msg)
        self.job_queue.put({'msg': msg, 'img': img})

    def run(self):
        while self.is_running():
            msg_dict = self.job_queue.get(block=True)
            if msg_dict is None:
                break

            msg, img = msg_dict['msg'], msg_dict['img']

            t1 = time.time()
            file_name = msg['file_name'] if 'file_name' in msg else ''

            self.infer_n_ims += 1
            self.infer_time1 += time.time() - t1
            # print('avg-infer time1: {}'.format(self.infer_time1 / self.infer_n_ims))

            t1 = time.time()
            # DO Object Detection
            results = self._detector(img, conf=self.confidence, iou=self.nms_thresh, imgsz=self.imgsz, verbose=False)
            res_msg = trans_det_results(results[0], self.kpt_links)
            res_msg['file_name'] = file_name
            res_msg['dataset'] = self.dataset_name
            if 'client_id' in msg:
                res_msg['client_id'] = msg['client_id']
            if 'file_name' in msg:
                res_msg['file_name'] = msg['file_name']
            if 'img_id' in msg:
                res_msg['img_id'] = msg['img_id']
            if 'img_total' in msg:
                res_msg['img_total'] = msg['img_total']
            res_msg['time_used'] = time.time() - t1

            self.infer_time2 += time.time() - t1
            # print('avg-infer time2: {}'.format(self.infer_time2 / self.infer_n_ims))

            t1 = time.time()

            self._result_writer.publish(res_msg)

            if 'img_total' in msg and self.launch_next_emit:
                next_msg = def_msg('std_msgs::Boolean')
                next_msg['data'] = True
                self._next_writer.publish(next_msg)

            if self.b_use_shm:
                msg = self._show_writer.cvimg2sms_mem(img)
            msg['spirecv_msgs::2DTargets'] = res_msg
            self._show_writer.publish(msg)

            self.infer_time3 += time.time() - t1
            # print('avg-infer time3: {}, totally {} imgs'.format(self.infer_time3 / self.infer_n_ims, self.infer_n_ims))
            # END

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name', '-j',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    # args = parser.parse_args()
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = YOLOv11PoseNode_Cuda(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.join()
